import numpy as np 
from PIL import Image
import cv2

class static_resize:
    def __init__(self, input_size):
        self.h = input_size[0]
        self.w = input_size[1]
        print(input_size)
        
    def __call__(self, img):
        if len(img.split()) == 3:
            padded_img = np.ones((self.h, self.w, 3), dtype=np.uint8) * 114
        else:
            padded_img = np.ones((self.h, self.w), dtype=np.int8) * 114
    
        r = min(self.h / img.size[0], self.w / img.size[1])
        resized_img = img.resize((int(r * img.size[0]), int(r * img.size[1])),
                            Image.Resampling.LANCZOS)
        padded_img[:int(img.size[1] * r), :int(img.size[0] * r)] = resized_img
        # cv2.namedWindow('hsv', 0)
        # cv2.imshow('hsv', padded_img)
        # cv2.waitKey(50)
        padded_img = Image.fromarray(padded_img)
        # padded_img.show()

        return padded_img
    
class augment_hsv:
    def __init__(self,):
        pass
    def __call__(self, img, hgain=5, sgain=30, vgain=30):
        img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
        hsv_augs = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain]  # random gains
        hsv_augs *= np.random.randint(0, 2, 3)  # random selection of h, s, v
        hsv_augs = hsv_augs.astype(np.int16)
        img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.int16)

        img_hsv[..., 0] = (img_hsv[..., 0] + hsv_augs[0]) % 180
        img_hsv[..., 1] = np.clip(img_hsv[..., 1] + hsv_augs[1], 0, 255)
        img_hsv[..., 2] = np.clip(img_hsv[..., 2] + hsv_augs[2], 0, 255)

        cv2.cvtColor(img_hsv.astype(img.dtype), cv2.COLOR_HSV2BGR, dst=img)  # no return needed
        
        # cv2.namedWindow('hsv', 0)
        # cv2.imshow('hsv', img)
        # cv2.waitKey(50)
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        return img